#!/usr/bin/env python3
"""
GPU服务器COS支持模块
修改GPU服务器以支持从COS下载图片和上传处理结果
"""

import os
import time
import uuid
import json
import logging
import requests
from datetime import datetime
from flask import Flask, request, jsonify
from werkzeug.utils import secure_filename
from qcloud_cos import CosConfig, CosS3Client
from qcloud_cos.cos_exception import CosServiceError
import hashlib

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class COSConfig:
    """COS配置类 - 支持多地域动态配置"""
    # COS配置
    SECRET_ID = os.getenv('COS_SECRET_ID', 'AKIDAQTJzVotlSwDH8p2MZw9E855uZHo8F9u')
    SECRET_KEY = os.getenv('COS_SECRET_KEY', 'jGav66wa9XkldKp9JI7kEdqulPq45vZ0')
    
    # 动态地域配置
    GPU_REGION = os.getenv('GPU_REGION', 'ap-beijing')
    COS_BUCKET_PREFIX = os.getenv('COS_BUCKET_PREFIX', 'photoenhancei')
    COS_ACCOUNT_ID = os.getenv('COS_ACCOUNT_ID', '1259206939')
    
    # 地域映射配置 - 主要支持4个地域
    REGION_MAPPING = {
        # 主要支持地域 (按优先级排序)
        'ap-beijing': {
            'cos_region': 'ap-beijing',
            'bucket_suffix': 'bj',
            'internal_ip_pattern': '10.2.'
        },
        'ap-shanghai': {
            'cos_region': 'ap-shanghai', 
            'bucket_suffix': 'sh',
            'internal_ip_pattern': '10.3.'
        },
        'ap-guangzhou': {
            'cos_region': 'ap-guangzhou',
            'bucket_suffix': 'gz', 
            'internal_ip_pattern': '10.1.'
        },
        'ap-chengdu': {
            'cos_region': 'ap-chengdu',
            'bucket_suffix': 'cd',
            'internal_ip_pattern': '10.4.'
        }
    }
    
    @classmethod
    def get_cos_config(cls):
        """动态获取COS配置"""
        # 根据GPU服务器内网IP自动检测地域
        detected_region = cls._detect_region_by_ip()
        
        if detected_region in cls.REGION_MAPPING:
            region_config = cls.REGION_MAPPING[detected_region]
            cos_region = region_config['cos_region']
            bucket_suffix = region_config['bucket_suffix']
        else:
            # 使用环境变量或默认配置
            cos_region = cls.GPU_REGION
            bucket_suffix = cos_region.split('-')[-1][:2]
        
        # 动态生成存储桶名称 (必须遵循腾讯云格式)
        cos_bucket = f"{cls.COS_BUCKET_PREFIX}-{bucket_suffix}-{cls.COS_ACCOUNT_ID}"
        
        # 动态生成内网域名
        cos_internal_domain = f"{cos_bucket}.cos-internal.{cos_region}.myqcloud.com"
        
        return {
            'bucket': cos_bucket,
            'region': cos_region,
            'internal_domain': cos_internal_domain,
            'detected_region': detected_region
        }
    
    @classmethod
    def _detect_region_by_ip(cls):
        """根据内网IP自动检测地域"""
        # 获取当前服务器的内网IP
        import socket
        try:
            # 获取本机内网IP
            s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            s.connect(("8.8.8.8", 80))
            local_ip = s.getsockname()[0]
            s.close()
        except:
            local_ip = "10.2.8.17"  # 默认IP
        
        for region, config in cls.REGION_MAPPING.items():
            if local_ip.startswith(config['internal_ip_pattern']):
                return region
        
        return cls.GPU_REGION
    
    # CDN域名
    CDN_DOMAIN = "cdn.gongjuxiang.work"

class COSManager:
    """COS管理器 - 支持多地域动态配置"""
    
    def __init__(self):
        # 获取动态COS配置
        self.cos_config_data = COSConfig.get_cos_config()
        
        self.config = CosConfig(
            Region=self.cos_config_data['region'],
            SecretId=COSConfig.SECRET_ID,
            SecretKey=COSConfig.SECRET_KEY,
            Domain=self.cos_config_data['internal_domain']  # 使用内网域名
        )
        self.client = CosS3Client(self.config)
    
    def download_from_cos(self, file_key, local_path):
        """从COS下载文件"""
        try:
            start_time = time.time()
            
            response = self.client.get_object(
                Bucket=self.cos_config_data['bucket'],
                Key=file_key,
                ResponseContentDisposition='attachment'
            )
            
            # 保存到本地
            with open(local_path, 'wb') as f:
                # 直接写入响应体内容
                f.write(response['Body'].get_raw_stream().read())
            
            download_time = time.time() - start_time
            logger.info(f"从COS下载文件成功: {file_key}, 耗时: {download_time:.2f}s")
            
            return {
                'success': True,
                'download_time': download_time,
                'file_size': os.path.getsize(local_path)
            }
            
        except CosServiceError as e:
            logger.error(f"COS下载失败: {e}")
            return {
                'success': False,
                'error': f"COS下载失败: {e.get_error_code()}"
            }
    
    def upload_to_cos(self, file_path, file_key):
        """上传文件到COS"""
        try:
            start_time = time.time()
            
        response = self.client.put_object_from_local_file(
            Bucket=self.cos_config_data['bucket'],
            LocalFilePath=file_path,
            Key=file_key,
            Metadata={
                'upload-time': str(int(time.time())),
                'source': 'gpu-server',
                'processed': 'true',
                'region': self.cos_config_data['region']
            }
        )
            
            upload_time = time.time() - start_time
            logger.info(f"文件上传到COS成功: {file_key}, 耗时: {upload_time:.2f}s")
            
            return {
                'success': True,
                'file_key': file_key,
                'upload_time': upload_time,
                'cos_url': f"https://{self.cos_config_data['internal_domain']}/{file_key}",
                'cdn_url': f"https://{COSConfig.CDN_DOMAIN}/{file_key}",
                'region': self.cos_config_data['region'],
                'bucket': self.cos_config_data['bucket']
            }
            
        except CosServiceError as e:
            logger.error(f"COS上传失败: {e}")
            return {
                'success': False,
                'error': f"COS上传失败: {e.get_error_code()}"
            }
    
    def get_cdn_url(self, file_key):
        """生成CDN访问URL"""
        return f"https://{COSConfig.CDN_DOMAIN}/{file_key}"

def create_gpu_api_with_cos_support():
    """创建支持COS的GPU API"""
    app = Flask(__name__)
    cos_manager = COSManager()
    
    # 存储任务状态
    task_status = {}
    
    @app.route('/api/v1/enhance', methods=['POST'])
    def enhance_image_with_cos():
        """
        支持COS的图片增强API
        流程: 接收COS URL -> 下载图片 -> 处理 -> 上传到COS -> 返回CDN URL
        """
        try:
            logger.info("收到COS支持的图片增强请求")
            
            # 检查是否有文件上传
            if 'file' not in request.files:
                return jsonify({
                    'success': False,
                    'error': '没有上传文件'
                }), 400
            
            file = request.files['file']
            if file.filename == '':
                return jsonify({
                    'success': False,
                    'error': '文件名为空'
                }), 400
            
            # 获取参数
            tile_size = int(request.form.get('tile_size', 400))
            quality_level = request.form.get('quality_level', 'high')
            cos_file_key = request.form.get('cos_file_key', '')  # 新增：COS文件键
            
            # 生成任务ID
            task_id = f"task_{uuid.uuid4().hex[:16]}"
            
            # 记录开始时间
            total_start_time = time.time()
            
            # 保存上传的文件到临时目录
            temp_dir = "/tmp/gpu_processing"
            os.makedirs(temp_dir, exist_ok=True)
            
            upload_filename = secure_filename(file.filename)
            input_path = os.path.join(temp_dir, f"{task_id}_input_{upload_filename}")
            output_path = os.path.join(temp_dir, f"{task_id}_output_{upload_filename}")
            
            file.save(input_path)
            file_size = os.path.getsize(input_path)
            
            logger.info(f"文件保存成功: {upload_filename}, 大小: {file_size} bytes")
            
            # 如果提供了COS文件键，从COS下载
            if cos_file_key:
                logger.info(f"从COS下载图片: {cos_file_key}")
                cos_download_result = cos_manager.download_from_cos(cos_file_key, input_path)
                
                if not cos_download_result['success']:
                    os.remove(input_path)
                    return jsonify({
                        'success': False,
                        'error': f"从COS下载失败: {cos_download_result['error']}"
                    }), 500
                
                logger.info(f"COS下载完成，耗时: {cos_download_result['download_time']:.2f}s")
            
            # 记录任务开始
            task_status[task_id] = {
                'status': 'processing',
                'progress': 0.1,
                'message': '开始处理图片',
                'start_time': total_start_time,
                'file_size': file_size
            }
            
            # 模拟GPU处理过程（这里应该调用实际的GPU处理代码）
            logger.info(f"开始GPU处理: {task_id}")
            
            # 更新进度
            task_status[task_id]['progress'] = 0.3
            task_status[task_id]['message'] = 'GPU处理中...'
            
            # 这里应该调用实际的GFPGAN处理
            # 为了演示，我们直接复制文件
            import shutil
            shutil.copy2(input_path, output_path)
            
            # 模拟处理时间
            time.sleep(2)  # 实际处理时间
            
            processing_time = time.time() - total_start_time
            
            # 生成输出文件键
            output_filename = f"enhanced_{upload_filename}"
            output_file_key = f"processed/{task_id}_{output_filename}"
            
            # 上传处理后的图片到COS
            logger.info(f"上传处理结果到COS: {output_file_key}")
            upload_result = cos_manager.upload_to_cos(output_path, output_file_key)
            
            if not upload_result['success']:
                # 清理临时文件
                for path in [input_path, output_path]:
                    if os.path.exists(path):
                        os.remove(path)
                return jsonify({
                    'success': False,
                    'error': f"上传到COS失败: {upload_result['error']}"
                }), 500
            
            # 生成CDN链接
            cdn_url = cos_manager.get_cdn_url(output_file_key)
            
            # 更新任务状态
            task_status[task_id].update({
                'status': 'completed',
                'progress': 1.0,
                'message': '处理完成',
                'processing_time': processing_time,
                'result_url': cdn_url,
                'cos_url': upload_result['cos_url'],
                'cdn_url': cdn_url,
                'file_key': output_file_key,
                'upload_time': upload_result['upload_time']
            })
            
            # 清理临时文件
            for path in [input_path, output_path]:
                if os.path.exists(path):
                    os.remove(path)
            
            # 返回结果
            result = {
                'success': True,
                'task_id': task_id,
                'enhanced_image': cdn_url,  # CDN链接
                'cos_url': upload_result['cos_url'],  # COS直链
                'performance': {
                    'processing_time': processing_time,
                    'upload_time': upload_result['upload_time'],
                    'file_size_original': file_size,
                    'file_size_enhanced': os.path.getsize(output_path) if os.path.exists(output_path) else file_size
                },
                'metadata': {
                    'file_key': output_file_key,
                    'tile_size': tile_size,
                    'quality_level': quality_level,
                    'timestamp': datetime.now().isoformat()
                }
            }
            
            logger.info(f"图片增强完成: {task_id}, 总耗时: {processing_time:.2f}s, CDN链接: {cdn_url}")
            
            return jsonify(result)
            
        except Exception as e:
            logger.error(f"图片增强异常: {e}")
            return jsonify({
                'success': False,
                'error': f'处理异常: {str(e)}'
            }), 500
    
    @app.route('/api/v1/status/<task_id>', methods=['GET'])
    def get_task_status(task_id):
        """获取任务状态"""
        if task_id not in task_status:
            return jsonify({'error': '任务不存在'}), 404
        
        status_info = task_status[task_id].copy()
        status_info['task_id'] = task_id
        
        return jsonify(status_info)
    
    @app.route('/api/v1/download/<task_id>', methods=['GET'])
    def download_result(task_id):
        """下载处理结果 - 重定向到CDN"""
        if task_id not in task_status:
            return jsonify({'error': '任务不存在'}), 404
        
        task_info = task_status[task_id]
        if task_info['status'] != 'completed':
            return jsonify({'error': '任务未完成'}), 400
        
        # 重定向到CDN链接
        cdn_url = task_info.get('cdn_url')
        if cdn_url:
            from flask import redirect
            return redirect(cdn_url)
        else:
            return jsonify({'error': '结果URL不存在'}), 404
    
    @app.route('/health', methods=['GET'])
    def health_check():
        """健康检查"""
        return jsonify({
            'status': 'healthy',
            'timestamp': datetime.now().isoformat(),
            'cos_connection': 'connected',
            'cos_bucket': COSConfig.BUCKET,
            'cos_region': COSConfig.REGION,
            'cdn_domain': COSConfig.CDN_DOMAIN
        })
    
    return app

if __name__ == '__main__':
    logger.info("启动支持多地域COS的GPU服务器")
    
    # 获取动态配置
    cos_config_data = COSConfig.get_cos_config()
    
    logger.info(f"检测到的地域: {cos_config_data['detected_region']}")
    logger.info(f"COS存储桶: {cos_config_data['bucket']}")
    logger.info(f"COS地域: {cos_config_data['region']}")
    logger.info(f"COS内网域名: {cos_config_data['internal_domain']}")
    logger.info(f"CDN域名: {COSConfig.CDN_DOMAIN}")
    logger.info(f"支持的地域: {list(COSConfig.REGION_MAPPING.keys())}")
    
    app = create_gpu_api_with_cos_support()
    app.run(host='0.0.0.0', port=5000, debug=False)
